
#define GLEW_STATIC
#include "include\GL\glew.h"
#define FREEGLUT_STATIC
#define FREEGLUT_LIB_PRAGMAS 0
#include "include\GL\glut.h"

#include <stdio.h>
#include <math.h>
#include <time.h>

#include "include\vector.h"
#include "include\matrix.h"
#include "include\shader.h"
#include "include\FluidSolver.h"
#include "include\Fluid.h"

#define INDEX(i,j) ((i)+(sideLength_+2)*(j))

class FluidDynamics {
	shader colourShader_;
	GLint viewport_width_;
	GLint viewport_height_;
	mat4 modelToProjection_;

	// Fluid Simulation Variables
	Fluid fire;
	Fluid smoke;

	float *velocityX_;
	float *velocityY_;
	float *previousVelocityX_;
	float *previousVelocityY_;
	float viscosity_;

	int paintColour_;

	enum {RED, GREEN, BLUE};

	int sideLength_; // How many cells wide and high. (N)
	float deltaTime_; // The change in time
	float force_;
	float source_;

	int numberOfIndices_;

	int mouse_down_[3];
	int omx_, omy_, mx_, my_;

	float *vertices;
	float *colours;
	GLuint *indices;

	int flameX;
	int flameY;

	void getFromUI() {
		int size = (sideLength_ + 2) * (sideLength_ + 2);

		for (int i = 0; i != size; ++i) {
			previousVelocityX_[i] = 0.0f;
			previousVelocityY_[i] = 0.0f;

			fire.setPreviousDensity(i, 0.0f);
		}

		if ( !mouse_down_[0] && !mouse_down_[2] ) return;

		int i = (int)((mx_ /(float)viewport_width_) * sideLength_ + 1);
		int j = (int)(((viewport_height_ - my_) / (float)viewport_height_) * sideLength_ + 1);

		if (i < 1 || i > sideLength_ || j < 1 || j > sideLength_ ) return;

		if ( mouse_down_[0] ) {
			previousVelocityX_[INDEX(i,j)] = force_ * (mx_-omx_);
			previousVelocityY_[INDEX(i,j)] = force_ * (omy_-my_);
		}

		if ( mouse_down_[2] ) {
			flameX = i;
			flameY = j;
		}

		omx_ = mx_;
		omy_ = my_;
	} // getFromUI()

	void updateKey( unsigned char key, int x, int y )
	{
		switch ( key )
		{
		case 'q':
		case 'Q':
			source_ += 0.5;
			printf("source: %f\n", source_);
			break;

		case 'a':
		case 'A':
			source_ -= 0.5;
			printf("source: %f\n", source_);
			break;

		case 'w':
		case 'W':
			force_ += 0.005;
			printf("force: %f\n", force_);
			break;

		case 's':
		case 'S':
			force_ -= 0.005;
			printf("force: %f\n", force_);
			break;

		case 'e':
		case 'E':
			fire.setDiffusion(fire.getDiffusion() + 0.0001);
			printf("diffusion: %f\n", fire.getDiffusion());
			break;

		case 'd':
		case 'D':
			fire.setDiffusion(fire.getDiffusion() - 0.0001);
			if (fire.getDiffusion() < 0) fire.setDiffusion(0);
			printf("diffusion: %f\n", fire.getDiffusion());
			break;

		case 'r':
		case 'R':
			viscosity_ += 0.001;
			printf("viscosity: %f\n", viscosity_);
			break;

		case 'f':
		case 'F':
			viscosity_ -= 0.001;
			if (viscosity_ < 0) viscosity_ = 0;
			printf("viscosity: %f\n", viscosity_);
			break;

		case 'c':
		case 'C':
			clearData();
			break;
		}
	}

	void update() {
		getFromUI();
		
		srand(time(NULL));
		float num;

		num = (float)rand() / RAND_MAX * 1.5f;
		
		fire.setPreviousDensity(INDEX(flameX, flameY), source_ * num);
		fire.setPreviousDensity(INDEX(flameX, flameY + 1), source_ * num);
		fire.setPreviousDensity(INDEX(flameX, flameY - 1), source_ * num);
		fire.setPreviousDensity(INDEX(flameX - 1, flameY), source_ * num);
		fire.setPreviousDensity(INDEX(flameX + 1, flameY), source_ * num);

		int size = (sideLength_ + 2) * (sideLength_ + 2);

		for (int i = 0; i != size; ++i) {
			previousVelocityY_[i] += fire.getDensity()[i] * force_;
			
			num = (float)rand() / RAND_MAX;
			previousVelocityX_[i] += num * 0.20f;
			num = (float)rand() / RAND_MAX;
			previousVelocityX_[i] += -num * 0.20f;
		}

		FluidSolver::velocityStep(sideLength_, velocityX_, velocityY_, previousVelocityX_, previousVelocityY_, viscosity_, deltaTime_);

		fire.update(deltaTime_, velocityX_, velocityY_);
	} // update()

	void drawDensity() {
		
		colourShader_.render(modelToProjection_);
		
		int size = (sideLength_ + 2) * (sideLength_ + 2);

		int colour = 0;
		for (int i = 0; i != size; ++i) {
			colours[colour] = fire.getDensity()[i];
			colours[colour + 1] = fire.getDensity()[i] * 2.9f / 4.0f;
			colours[colour + 2] = fire.getDensity()[i] * 1.6f / 3.0f;
			colours[colour + 3] = fire.getDensity()[i];
			
			colour += 4;
		} // for

		glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 2*sizeof(float), (void*)vertices);
		glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 4*sizeof(float), (void*)colours);
		glEnableVertexAttribArray(0);
		glEnableVertexAttribArray(1);
		glDrawElements(GL_TRIANGLES, numberOfIndices_, GL_UNSIGNED_INT, indices);

	} // drawDensity()

	void build_camera_matrix()
	{
		mat4 cameraToWorld;
		cameraToWorld.loadIdentity();
		cameraToWorld.translate(0, 0, 10);

		mat4 worldToCamera;
		cameraToWorld.invertQuick(worldToCamera);

		mat4 cameraToProjection;
		cameraToProjection.loadIdentity();
		float n = 6, f = 30;
		cameraToProjection.frustum(-n, n, -n, n, n, f);

		modelToProjection_ = worldToCamera * cameraToProjection;
	}

	void render() {
		update();
		
		glClearColor(0.6f, 0.7f, 0.9f, 1);
		glViewport(0, 0, viewport_width_, viewport_height_);
		glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT);

		drawDensity();

		glutSwapBuffers();
	} // render()

	void allocateData() {
		int size = (sideLength_ + 2) * (sideLength_ + 2);

		velocityX_ = (float*)malloc(size * sizeof(float));
		velocityY_ = (float*)malloc(size * sizeof(float));
		previousVelocityX_ = (float*)malloc(size * sizeof(float));
		previousVelocityY_ = (float*)malloc(size * sizeof(float));

		vertices = (GLfloat*)malloc(size * 2 * sizeof(float));
		colours = (GLfloat*)malloc(size * 4 * sizeof(float));
		indices = (GLuint*)malloc(6 * (sideLength_ + 1) * (sideLength_ + 1) * sizeof(GLuint));

		printf("Number of vertices = %d\n", size);
		printf("Number of indices = %d\n\n", 6 * (sideLength_ + 1) * (sideLength_ + 1));
		numberOfIndices_ = 6 * (sideLength_ + 1) * (sideLength_ + 1);

		if (!vertices || !indices) {
			printf("Couldn't allocate necessary arrays! Exiting!\n");  
			exit(1);
		}
	} // allocateData()

	void clearData() {
		int size = (sideLength_ + 2) * (sideLength_ + 2);

		for (int i = 0; i != size; ++i) {
			velocityX_[i] = 0.0f;
			velocityY_[i] = 0.0f;
			previousVelocityX_[i] = 0.0f;
			previousVelocityY_[i] = 0.0f;
		}

		fire.clearData();
	} // ClearData()

	void createLattice() {
		int vertex = 0;

		float xGap = 20.0 / (sideLength_ + 1);
		float yGap = 20.0 / (sideLength_ + 1);

		int numberOfVertsWide = sideLength_ + 2;
		int numberOfVertsHigh = sideLength_ + 2;

		for (int j = 0; j != numberOfVertsHigh; ++j) {
			for (int i = 0; i != numberOfVertsWide; ++i) {
				vertices[vertex] = -10 + (i * xGap);
				vertices[vertex + 1] = -10 + (j * yGap);

				vertex += 2;
			} // for
		} // for

		int index = 0;

		int rows = sideLength_ + 1;
		int trianglePairsPerRow = sideLength_ + 1;

		for (int j = 0; j != rows; ++j) {
			for (int i = 0; i != trianglePairsPerRow; ++i) {

				indices[index] = i + (j * (sideLength_ + 2));
				indices[index + 1] = i + ((j + 1) * (sideLength_ + 2)) + 1;
				indices[index + 2] = i + ((j + 1) * (sideLength_ + 2));

				indices[index + 3] = i + (j * (sideLength_ + 2));
				indices[index + 4] = i + (j * (sideLength_ + 2)) + 1;
				indices[index + 5] = i + ((j + 1) * (sideLength_ + 2)) + 1;

				index += 6;
			} // for
		} // for

	} // createLattice()

	void initializeSimulation() {

		sideLength_ = 64;
		deltaTime_ = 0.1f;
		force_ = 1.71f;
		source_ = 90.0f;

		viscosity_ = 0.0f;

		allocateData();
		createLattice();

		fire.initialize(sideLength_, 0.0041f); // 
		//smoke.initialize(sideLength_, 0.0001f);

		flameX = sideLength_/2;
		flameY = 20;

		clearData();
	} // initializeSimulation()

	FluidDynamics() {
		colourShader_.init(
			"attribute vec4 position;"
			"attribute vec4 vertColour;"
			"varying vec4 colour;"
			"uniform mat4 modelToProjection;"

			"void main() {"
			"  colour = vertColour;"
			"  gl_Position = modelToProjection * position;"
			"}",

			"varying vec4 colour;"
			"void main() {"
			"  gl_FragColor = colour;"
			"}"
		);

		build_camera_matrix();

		initializeSimulation();

		//paintColour_ = RED;

		glEnable(GL_BLEND);
		glBlendFunc(GL_SRC_COLOR, GL_ONE_MINUS_SRC_COLOR) ;

		// Display hotkeys
		printf("Increase/Decrease:\n");
		printf("Source input (initial flame intensity): <Q/A>\n");
		printf("Force (hoe much the heat influences lift): <W/S>\n");
		printf("Diffusion (of the heat into the environment): <E/D>\n");
		printf("Viscosity: <R/F>\n");

	} // FluidDynamics()

	void set_viewport(int w, int h) {
		viewport_width_ = w;
		viewport_height_ = h;
	} // setViewport()

	void updateMouse(int button, int state, int x, int y) {
		omx_ = mx_ = x;
		omy_ = my_ = y;

		mouse_down_[button] = state == GLUT_DOWN;
	} // updateMouse()

	void updateMouseMotion(int x, int y) {
		mx_ = x;
		my_ = y;
	} // updateMouseMotion()

public:
	static FluidDynamics &get()
	{
		static FluidDynamics singleton;
		return singleton;
	}

	// interface from GLUT
	static void reshape(int w, int h) { get().set_viewport(w, h); }
	static void display() { get().render(); }
	static void timer(int value) { glutTimerFunc(30, timer, 1); glutPostRedisplay(); }
	static void mouse(int button, int state, int x, int y) {get().updateMouse(button, state, x, y);}
	static void mouseMotion(int x, int y) {get().updateMouseMotion(x, y);}
	static void key( unsigned char key, int x, int y) {get().updateKey(key, x, y); }
};

int main(int argc, char **argv)
{
	glutInit(&argc, argv);
	glutInitDisplayMode(GLUT_RGBA|GLUT_DEPTH|GLUT_DOUBLE);
	glutInitWindowSize(500, 500);
	glutCreateWindow("Fluid Dynamics Demo - Andrew Dyer");
	glewInit();
	if (!glewIsSupported("GL_VERSION_2_0") )
	{
		printf("OpenGL 2 is required!\n");
		return 1;
	}
	glutDisplayFunc(FluidDynamics::display);
	glutReshapeFunc(FluidDynamics::reshape);
	glutKeyboardFunc(FluidDynamics::key);
	glutMouseFunc(FluidDynamics::mouse);
	glutMotionFunc(FluidDynamics::mouseMotion);
	glutTimerFunc(33, FluidDynamics::timer, 1);
	glutMainLoop();
	return 0;
}

